"""
Ollama embedder implementation.
"""

import asyncio
import json
from typing import List, Optional
from .base import IEmbedder, EmbeddingResponse, EmbedderInfo, DocumentType, TaskType

try:
    import aiohttp
except ImportError:
    raise ImportError("aiohttp package is required for Ollama embedder. Install with: pip install aiohttp")

# Constants
MAX_BATCH_RETRIES = 3
INITIAL_RETRY_DELAY_MS = 1000


class InstructionManager:
    """Manages task-specific instructions for embeddings."""
    
    @staticmethod
    def get_instruction(task: TaskType, doc_type: DocumentType, text: str) -> str:
        """Get task and document type specific instruction."""
        if task == "retrieval_query":
            if doc_type == "code":
                return f"Represent this code search query for retrieving relevant code snippets: {text}"
            else:
                return f"Represent this search query for retrieving relevant documents: {text}"
        else:  # retrieval_document
            if doc_type == "code":
                return f"Represent this code snippet for search: {text}"
            else:
                return f"Represent this document for search: {text}"


class OllamaEmbedder(IEmbedder):
    """Ollama embedder implementation."""
    
    def __init__(self, base_url: str, model_id: str = "bge-m3:latest"):
        self.base_url = base_url.rstrip('/')
        self.model_id = model_id
    
    async def create_embeddings(
        self, 
        texts: List[str], 
        doc_type: DocumentType = "text", 
        task: TaskType = "retrieval_document"
    ) -> EmbeddingResponse:
        """Create embeddings using Ollama API."""
        embeddings = []
        
        for text in texts:
            instructed_text = InstructionManager.get_instruction(task, doc_type, text)
            embedding = await self._embed_with_retries(instructed_text)
            embeddings.append(embedding)
        
        return EmbeddingResponse(embeddings=embeddings)
    
    async def _embed_with_retries(self, text: str) -> List[float]:
        """Embed text with retry logic."""
        for attempt in range(MAX_BATCH_RETRIES):
            try:
                async with aiohttp.ClientSession() as session:
                    async with session.post(
                        f"{self.base_url}/api/embeddings",
                        json={
                            "model": self.model_id,
                            "prompt": text
                        },
                        headers={"Content-Type": "application/json"}
                    ) as response:
                        if not response.ok:
                            raise Exception(f"Ollama API error: {response.status} {response.reason}")
                        
                        data = await response.json()
                        
                        if not data.get("embedding") or not isinstance(data["embedding"], list):
                            raise Exception("Invalid response format from Ollama API")
                        
                        return data["embedding"]
                        
            except Exception as error:
                if attempt < MAX_BATCH_RETRIES - 1:
                    delay_ms = INITIAL_RETRY_DELAY_MS * (2 ** attempt)
                    await asyncio.sleep(delay_ms / 1000.0)
                else:
                    raise RuntimeError(f"Failed to create embedding after {MAX_BATCH_RETRIES} attempts: {error}")
        
        raise RuntimeError(f"Failed to create embedding after {MAX_BATCH_RETRIES} attempts")
    
    async def validate_configuration(self) -> tuple[bool, Optional[str]]:
        """Validate Ollama embedder configuration."""
        try:
            # Check if Ollama server is running
            async with aiohttp.ClientSession() as session:
                async with session.get(f"{self.base_url}/api/tags") as response:
                    if not response.ok:
                        return False, f"Cannot connect to Ollama server: {response.status} {response.reason}"
                    
                    data = await response.json()
                    models = data.get("models", [])
                    
                    # Check if model exists
                    model_exists = any(
                        model.get("name", "").startswith(self.model_id) 
                        for model in models
                    )
                    
                    if not model_exists:
                        return False, f"Model '{self.model_id}' not found. Available models: {[m.get('name') for m in models]}"
            
            # Test embedding creation
            test_embedding = await self._embed_with_retries("test")
            if not test_embedding or len(test_embedding) == 0:
                return False, "Embedding test returned empty result"
            
            return True, None
            
        except Exception as error:
            return False, f"Ollama validation failed: {error}"
    
    @property
    def embedder_info(self) -> EmbedderInfo:
        return EmbedderInfo(name="ollama")